import paddle
import paddle.nn as nn
import numpy as np
class Cbl(nn.Layer):
    def __init__(self,inputchannel,outputchannels,k=1,s=1,p=1):
        super(Cbl, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2D(inputchannel,outputchannels,kernel_size=k,stride=s,padding=p),
            nn.BatchNorm2D(outputchannels),
            nn.ReLU()
        )
    def forward(self,x):
        return self.conv(x)
class Resunit(nn.Layer):
    def __init__(self,inputchannel):
        super(Resunit, self).__init__()
        self.cbl1 = Cbl(inputchannel,inputchannel//2,p=0)
        self.cbl2 = Cbl(inputchannel//2,inputchannel,k=3)
    def forward(self,x):
        x1 = self.cbl1(x)
        x1 = self.cbl2(x1)
        x1 = x+x1
        return x1
class ResX(nn.Layer):
    def __init__(self,numx,inputchannel):
        super(ResX, self).__init__()
        self.cbl = Cbl(inputchannel,inputchannel*2,k=3,s=2)
        self.ResU = [Resunit(inputchannel*2) for i in range(numx)]
        self.numx = numx
    def forward(self,x):
        x = self.cbl(x)
        for i in range(self.numx):
            x = self.ResU[i](x)
        return x
class yoloblock(nn.Layer):
    def __init__(self,inputchannel,ouputchannel):
        super(yoloblock, self).__init__()
        self.conv1 = Cbl(inputchannel,ouputchannel,p=0)
        self.conv2 = Cbl(ouputchannel,ouputchannel*2,k=3)
        self.conv3 = Cbl(ouputchannel*2,ouputchannel,p=0)
        self.conv4 = Cbl(ouputchannel,ouputchannel*2,k=3)
        self.conv5 = Cbl(ouputchannel*2,ouputchannel,p=0)
    def forward(self,x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        return x

class Backbone(nn.Layer):
    def __init__(self):
        super(Backbone, self).__init__()
        self.cbl = Cbl(3,32,3)
        self.res1_1 = ResX(1,32)
        self.res2_2 = ResX(2,64)
        self.res3_8 = ResX(8,128)
        self.res4_8 = ResX(8,256)
        self.res5_4 = ResX(4,512)
    def forward(self,x):
        x = self.cbl(x)
        x = self.res1_1(x)
        x = self.res2_2(x)
        x = self.res3_8(x)
        x1 = self.res4_8(x)
        x2 = self.res5_4(x1)
        return x,x1,x2

class Neck(nn.Layer):
    def __init__(self):
        super(Neck, self).__init__()
        self.cbl1 = yoloblock(1024,512)
        self.cbl2 = yoloblock(768,256)
        self.cbl3 = yoloblock(384,128)
        self.cbl4 = Cbl(512,256,p=0)
        self.cbl5 = Cbl(256,128,p=0)
        self.up1 = nn.UpsamplingBilinear2D(scale_factor=2)
        self.up2 = nn.UpsamplingBilinear2D(scale_factor=2)
    def forward(self,x1,x2,x3):
        x1 = self.cbl1(x1)
        x2 = paddle.concat([self.up1(self.cbl4(x1)),x2],axis=1)
        x2 = self.cbl2(x2)
        x3 = paddle.concat([self.up2(self.cbl5(x2)),x3],axis=1)
        x3 = self.cbl3(x3)
        return x1,x2,x3
class yolov3(nn.Layer):
    def __init__(self):
        super(yolov3, self).__init__()
        self.backbone = Backbone()
        self.neck = Neck()
        self.cbl1 = Cbl(512,1024,k=3)
        self.conv1 = nn.Conv2D(1024,3*(5+4),kernel_size=1)
        self.cbl2 = Cbl(256,512,k=3)
        self.conv2 = nn.Conv2D(512,3*(5+4),kernel_size=1)
        self.cbl3 = Cbl(128,256,k=3)
        self.conv3 = nn.Conv2D(256,3*(5+4),kernel_size=1)
    def forward(self,x):
        x3,x2,x1 = self.backbone(x)
        x1,x2,x3 = self.neck(x1,x2,x3)
        x1 = self.conv1(self.cbl1(x1))
        x2 = self.conv2(self.cbl2(x2))
        x3 = self.conv3(self.cbl3(x3))
        return x1,x2,x3

if __name__ == '__main__':
    data = np.zeros((1, 3, 416, 416))
    data = paddle.to_tensor(data)
    data = data.astype(paddle.float32)
    model = yolov3()
    model.eval()
    out1,out2,out3 = model(data)
    print(out1.shape)
    print(out2.shape)
    print(out3.shape)







